"""
更新内容：1.增加参数选择
1.1 输入需要跑分析的物种名文件，如下格式：
Aal2
Aan1
Aar4
1.2 输入目标路径名，即拷贝物种相关文件到的目录名
2. 原来的需要拷贝的_mRNA_countsLength.txt文件，现在不拷贝该文件，
通过bed文件和chrom.sizes.txt文件手动计算物种的基因组长度和基因数量
"""

"""
    拷贝IMP数据库中每个物种的相关文件到一个目录下面
        每个物种是三个文件：
            pfam.table.txt                          # pfam的注释结果文件
            .bed                                    # bed文件
            eggnog.emapper.annotations.all.txt      # eggnog的注释结果
"""


import os
import shutil
import subprocess
import argparse

parser = argparse.ArgumentParser()
parser.add_argument("--species", required=True)
parser.add_argument("--input", required=True)
parser.add_argument("--output", required=True)
args = parser.parse_args()

def whole_file_path(dir):
    """
        对一个文件夹下面的文件都拼成全路径
    """
    file_whole_path = []
    for root, dirs, files in os.walk(dir):
        for file in files:
            file_path = os.path.join(root, file)
            file_path = file_path.replace("\\", "/")
            file_whole_path.append(file_path)

    return file_whole_path


if __name__ == "__main__":
    #dir = "/data/genome/IMP"    # 数据所在路径
    dir = args.input
    print(dir)
    # 获取一个目录下所有文件的全路径
    #whole_paths = whole_file_path(dir)
    # 提取有pfam结果的路径
    #pfam_paths = [x for x in whole_paths if x.endswith("pfam.table.txt")]

    #print("路径过滤完成！")

    ## 读取需要跑分析的物种名文件
    with open(args.species, "r") as file:
        lines = file.readlines()
        lines = [line.strip() for line in lines]


    for plant_name in lines:
        # 如果需要的三个文件路径总是固定的，就直接去判断这三个文件存不存在
        pfam_path = f"{dir}/{plant_name}/result/pfam/pfam.table.txt"
        bed_path = f"{dir}/{plant_name}/result/{plant_name}.bed"
        eggnog_path = f"{dir}/{plant_name}/result/eggnog/eggnog.emapper.annotations.all.txt"
        ##count_path = f"/data/genome/IMP/{plant_name}/{plant_name}_mRNA_countsLength.txt"
        
        # 如果三种文件都存在的物种，那么才拷贝到目的文件夹中
        if os.path.exists(bed_path) and os.path.exists(pfam_path) and os.path.exists(eggnog_path):
                print(pfam_path)
                # 根据bed文件和chrom.sizes.txt计算物种基因组长度和基因个数，并生成_mRNA_countsLength.txt文件到目的文件夹
                chrom_sizes_file = f"{dir}/{plant_name}/result/igv/{plant_name}.chrom.sizes.txt"
                out_path = f"{args.output}/{plant_name}/{plant_name}_mRNA_countsLength.txt"
                
            # 目的文件夹
                result_path = args.output + "/" + plant_name
                print (result_path)
                os.makedirs(result_path, exist_ok=True)
                shutil.copy(bed_path, result_path)
                shutil.copy(pfam_path, result_path)
                shutil.copy(eggnog_path, result_path)
                # 取bed文件第五列去重得到的个数记为基因个数
                command1 = f"awk -F '\t' '{{print $5}}' '{bed_path}' | sort | uniq | wc -l"
                result1 = subprocess.run(command1,shell=True,capture_output=True,text=True)
                gene_num = result1.stdout.strip()
                # 取chrom.sizes.txt文件第二列总和记为基因组长度
                command2 = f"awk -F '\t' '{{sum+=$2}}END{{print sum}}' '{chrom_sizes_file}'"
                result2 = subprocess.run(command2,shell=True,capture_output=True,text=True)
                genome_num = result2.stdout.strip()
                file = plant_name + "\t" + gene_num + "\t" + genome_num + "\n"
                #print(file)
                with open(out_path, 'w') as f:
                    f.write(file)
                #shutil.copy(count_path, result_path)

                print("   " + plant_name + "   已经复制完成! ")

    print("\n-----------复制完成-----------\n")

